import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from numpy.polynomial.chebyshev import Chebyshev

# 配置中文字体
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置中文字体为黑体
plt.rcParams["axes.unicode_minus"] = False  # 解决负号'-'显示为方块的问题


# 定义待逼近的复杂函数
def f(x):
    return np.sin(2 * np.pi * x) + 0.5 * np.cos(3 * np.pi * x) + x**2


# 基函数定义
def poly_basis(i, x):
    return x**i


def sin_basis(i, x):
    if i == 0:
        return np.sin(np.pi * x)
    else:
        return np.sin((i + 1) * np.pi * x)


def chebyshev_basis(i, x):
    return Chebyshev.basis(i)(x)


# 内积计算
def inner_product(phi_i, phi_j):
    result, _ = quad(lambda x: phi_i(x) * phi_j(x), 0, 1)
    return result


def inner_product_f_phi(f, phi_i):
    result, _ = quad(lambda x: f(x) * phi_i(x), 0, 1)
    return result


# 逼近函数计算
def galerkin_approximation(basis_functions, n):
    A = np.zeros((n, n))
    b = np.zeros(n)

    for i in range(n):
        for j in range(n):
            A[i, j] = inner_product(lambda x: basis_functions[i](x), lambda x: basis_functions[j](x))
        b[i] = inner_product_f_phi(f, lambda x: basis_functions[i](x))

    c = np.linalg.solve(A, b)
    # print(c)

    def u_approx(x):
        return sum(c[i] * basis_functions[i](x) for i in range(n))

    return u_approx


def VariousBasis():
    # 定义基函数数量
    basis_num = 4
    plot_pt_num = 100

    # 基函数集合
    poly_basis_functions = [lambda x, i=i: poly_basis(i, x) for i in range(basis_num)]
    sin_basis_functions = [lambda x, i=i: sin_basis(i, x) for i in range(basis_num)]
    chebyshev_basis_functions = [lambda x, i=i: chebyshev_basis(i, x) for i in range(basis_num)]

    # 计算逼近函数
    u_poly = galerkin_approximation(poly_basis_functions, basis_num)
    u_sin = galerkin_approximation(sin_basis_functions, basis_num)
    u_chebyshev = galerkin_approximation(chebyshev_basis_functions, basis_num)

    # 绘制解析解
    x_fine = np.linspace(0, 1, 1000)
    y_fine = f(x_fine)
    plt.plot(x_fine, y_fine, label="解析解 $f(x)=\sin(2\pi x) + 0.5\cos(3\pi x)$", color="red")

    # 绘制多项式逼近
    x_fine = np.linspace(0, 1, plot_pt_num)
    y_poly = np.array([u_poly(x) for x in x_fine])
    plt.plot(x_fine, y_poly, label="多项式基函数 $x^i$", linestyle="--")
    # 获取曲线的颜色
    line_color = plt.gca().lines[-1].get_color()
    plt.scatter(x_fine, [u_poly(node) for node in x_fine], color=line_color, zorder=5)

    # 绘制三角函数逼近
    y_sin = np.array([u_sin(x) for x in x_fine])
    plt.plot(x_fine, y_sin, label="Sine basis $\sin((i+1)\pi x)$", linestyle="-.")
    # 获取曲线的颜色
    line_color = plt.gca().lines[-1].get_color()
    plt.scatter(x_fine, [u_sin(node) for node in x_fine], color=line_color, zorder=5)

    # 绘制Chebyshev逼近
    y_chebyshev = np.array([u_chebyshev(x) for x in x_fine])
    plt.plot(x_fine, y_chebyshev, label="Chebyshev basis $T_i(x)$", linestyle=":")
    # 获取曲线的颜色
    line_color = plt.gca().lines[-1].get_color()
    plt.scatter(x_fine, [u_chebyshev(node) for node in x_fine], color=line_color, zorder=5)

    # 绘图
    plt.legend()
    # 添加图名
    plt.title("不同基函数的逼近效果", fontsize=20)
    plt.xlabel("$x$")
    plt.ylabel("$f(x)$")
    plt.grid(True)
    plt.show()


def VariousPlotPtNum():
    # 定义基函数数量
    basis_num = 4

    # 基函数集合
    basis_functions = [lambda x, i=i: poly_basis(i, x) for i in range(basis_num)]

    # 计算逼近函数
    appro_curve = galerkin_approximation(basis_functions, basis_num)

    # 绘制解析解
    x_fine = np.linspace(0, 1, 1000)
    y_fine = f(x_fine)
    plt.plot(x_fine, y_fine, label="解析解 $f(x)=\sin(2\pi x) + 0.5\cos(3\pi x)$", color="red")

    for plot_pt_num in range(2, 10, 2):
        # 绘制多项式逼近
        x_fine = np.linspace(0, 1, plot_pt_num)
        y_poly = np.array([appro_curve(x) for x in x_fine])
        plt.plot(x_fine, y_poly, label=f"节点数{plot_pt_num}", linestyle="--")
        # 获取曲线的颜色
        line_color = plt.gca().lines[-1].get_color()
        plt.scatter(x_fine, [appro_curve(node) for node in x_fine], color=line_color, zorder=5)

    # 绘图
    plt.legend()
    # 添加图名
    plt.title("分片逼近的效果", fontsize=20)
    plt.xlabel("$x$")
    plt.ylabel("$f(x)$")
    plt.grid(True)
    plt.show()


def VariousBasisNum():
    # 定义基函数数量
    plot_pt_num = 10

    # 绘制解析解
    x_fine = np.linspace(0, 1, 1000)
    y_fine = f(x_fine)
    plt.plot(x_fine, y_fine, label="解析解 $f(x)=\sin(2\pi x) + 0.5\cos(3\pi x)$", color="red")

    for basis_num in range(2, 10, 2):

        # 基函数集合
        sin_basis_functions = [lambda x, i=i: sin_basis(i, x) for i in range(basis_num)]

        # 计算逼近函数
        u_sin = galerkin_approximation(sin_basis_functions, basis_num)

        # 绘制多项式逼近
        x_fine = np.linspace(0, 1, plot_pt_num)
        y_poly = np.array([u_sin(x) for x in x_fine])
        plt.plot(x_fine, y_poly, label=f"基函数数量{basis_num}", linestyle="--")
        # 获取曲线的颜色
        line_color = plt.gca().lines[-1].get_color()
        plt.scatter(x_fine, [u_sin(node) for node in x_fine], color=line_color, zorder=5)

    # 绘图
    plt.legend()
    # 添加图名
    plt.title("不同基函数数量的逼近效果", fontsize=20)
    plt.xlabel("$x$")
    plt.ylabel("$f(x)$")
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    # VariousBasis()
    VariousBasisNum()
    # VariousPlotPtNum()
